import os
import time
import glob
from scapy.all import rdpcap
from scapy.layers.inet import IP, TCP, UDP
from scapy.layers.l2 import Ether
from config.settings import get_config
from base.base_page import BasePage

# 配置参数
class Pcap_data_extraction(BasePage):
    def __init__(self, driver, device_type, logger=None):
        super().__init__(driver, device_type, logger=logger)  # 传递日志实例
        self.driver = driver
        self.config = get_config(device_type)  # 获取全局配置

    def wait_for_new_pcap(self):
        """等待新PCAP文件下载完成"""
        start_time = time.time()
        existing_files = set(glob.glob(os.path.join(self.config.download_dir, "*.pcap")))

        while True:
            current_files = set(glob.glob(os.path.join(self.config.download_dir, "*.pcap")))
            new_files = list(current_files - existing_files)

            if new_files:
                newest_file = max(new_files, key=os.path.getctime)
                # 验证文件是否完成下载
                initial_size = os.path.getsize(newest_file)
                time.sleep(1)
                if os.path.getsize(newest_file) == initial_size:
                    self.logger.info("文件已下载完成")
                    return newest_file

            if time.time() - start_time > 30:
                raise TimeoutError("等待PCAP文件下载超时")

            time.sleep(1)
    def get_newest_pcap(self, existing_files=None):
        """获取最新下载的PCAP文件"""
        pattern = os.path.join(self.config.download_dir, "*.pcap")
        all_files = glob.glob(pattern)
        if existing_files:
            new_files = list(set(all_files) - set(existing_files))
            if new_files:
                return max(new_files, key=os.path.getctime)
            self.logger.info("获取最新下载的PCAP文件成功")
        return max(all_files, key=os.path.getctime) if all_files else None

    def parse_pcap_with_scapy(self, download_file, required_fields):
        """
        动态解析PCAP文件
        :param download_file: PCAP文件路径
        :param required_fields: 需要提取的字段列表，可选值：
            ['src_ip', 'dst_ip', 'src_mac', 'dst_mac',
             'src_port', 'dst_port', 'protocol', 'timestamp']
        :return: 包含字典的列表，每个字典对应一个数据包
        """
        try:
            packets = rdpcap(download_file)
            results = []

            for pkt in packets:
                packet_info = {field: "N/A" for field in required_fields}

                # 解析两层信息
                if Ether in pkt:
                    if 'src_mac' in required_fields:
                        packet_info['src_mac'] = pkt[Ether].src
                    if 'dst_mac' in required_fields:
                        packet_info['dst_mac'] = pkt[Ether].dst

                # 解析三层信息
                if IP in pkt:
                    ip = pkt[IP]
                    if 'src_ip' in required_fields:
                        packet_info['src_ip'] = ip.src
                    if 'dst_ip' in required_fields:
                        packet_info['dst_ip'] = ip.dst
                    if 'protocol' in required_fields:
                        protocol_map = {6: 'TCP', 17: 'UDP', 1: 'ICMP'}
                        packet_info['protocol'] = protocol_map.get(ip.proto, str(ip.proto))

                # 解析四层信息
                if TCP in pkt:
                    tcp = pkt[TCP]
                    if 'src_port' in required_fields:
                        packet_info['src_port'] = tcp.sport
                    if 'dst_port' in required_fields:
                        packet_info['dst_port'] = tcp.dport
                elif UDP in pkt:
                    udp = pkt[UDP]
                    if 'src_port' in required_fields:
                        packet_info['src_port'] = udp.sport
                    if 'dst_port' in required_fields:
                        packet_info['dst_port'] = udp.dport

                # 时间戳处理
                if 'timestamp' in required_fields and hasattr(pkt, 'time'):
                    packet_info['timestamp'] = pkt.time

                results.append(packet_info)
                self.logger.debug(f"解析数据包成功: {packet_info}")

            self.logger.info(f"成功解析 {len(results)} 个数据包")
            return results

        except Exception as e:
            self.logger.error(f"PCAP解析失败: {str(e)}")
            return None